import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
import dgl.function as fn

# 随机生成图
num_nodes = 10
num_edges = 20
in_feats = 5
out_feats = 3
edge_feats = 4
# edges_src,edges_dst  1*num_edges,可能存在自连边
edges_src = torch.randint(0, num_nodes, (num_edges,))
edges_dst = torch.randint(0, num_nodes, (num_edges,))
g = dgl.graph((edges_src, edges_dst))
g = dgl.add_self_loop(g)

# 节点特征h，形状为 [num_nodes, in_feats]
g.ndata['h'] = torch.randn(num_nodes, in_feats)
# 边特征num_edges*edge_feats=30*4
g.edata['feat'] = torch.randn(g.num_edges(), edge_feats)


# 自定义边特征变换层
class EdgeUpdateLayer(nn.Module):
    def __init__(self, edge_feats, out_feats):
        super().__init__()
        # 输入维度映射到输出维度 4-》3
        self.edge_fc = nn.Linear(edge_feats, out_feats)

    def forward(self, g):
        # 对每条边的特征做线性变换
        g.edata['feat_new'] = self.edge_fc(g.edata['feat'])
        return g.edata['feat_new']


# 实例化并调用
edge_layer = EdgeUpdateLayer(edge_feats, 6)
new_edge_feats = edge_layer(g)
print('新边特征形状:', new_edge_feats.shape)





# 自定义GAT层，显式消息传递
class SimpleGATLayer(nn.Module):
    def __init__(self, in_feats, out_feats):
        super().__init__()
        self.fc = nn.Linear(in_feats, out_feats, bias=False)
        self.attn_fc = nn.Linear(2 * out_feats, 1, bias=False)

    def edge_attention(self, edges):
        z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)
        a = self.attn_fc(z2)
        return {'e': F.leaky_relu(a)}

    def message_func(self, edges):
        return {'z': edges.src['z'], 'e': edges.data['e']}

    def reduce_func(self, nodes):
        alpha = F.softmax(nodes.mailbox['e'], dim=1)
        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)
        return {'h_new': h}

    def forward(self, g, h):
        z = self.fc(h)
        g.ndata['z'] = z
        g.apply_edges(self.edge_attention)
        g.update_all(self.message_func, self.reduce_func)
        return g.ndata.pop('h_new')

# 简单两层GAT
class SimpleGAT(nn.Module):
    def __init__(self, in_feats, hidden_feats, out_feats):
        super().__init__()
        self.layer1 = SimpleGATLayer(in_feats, hidden_feats)
        self.layer2 = SimpleGATLayer(hidden_feats, out_feats)

    def forward(self, g, h):
        h = self.layer1(g, h)
        h = F.elu(h)
        h = self.layer2(g, h)
        return h

# 实例化并前向传播
model = SimpleGAT(in_feats, 8, out_feats)
# model输出的是节点特征 node_num*out_feats  10*3
out = model(g, g.ndata['h'])
print('输出特征:', out.shape)